#include "iocp_server.h"
#include <winsock2.h>
#include <windows.h>
#include <stdio.h>
#include <stdlib.h>

#pragma comment(lib, "ws2_32.lib")

#define PORT 8080
#define BUFFER_SIZE 1024
#define MAX_WORKER_THREADS 2

typedef struct {
    OVERLAPPED overlapped;
    SOCKET socket;
    char buffer[BUFFER_SIZE];
    WSABUF wsabuf;
    DWORD operation;
} PER_IO_DATA, *LPPER_IO_DATA;

HANDLE g_hIOCP = INVALID_HANDLE_VALUE;
BOOL g_bShutdown = FALSE;

// 定义回调函数类型
typedef struct{
    DataCallback callback; 
} WorkerThreadParam;

 
WorkerThreadParam params = {NULL};
DWORD WINAPI WorkerThread(LPVOID lpParameter);
void CleanupSocket(LPPER_IO_DATA pIoData);

 

void start_server(DataCallback callback) {
    WSADATA wsaData;
    SOCKET listenSocket = INVALID_SOCKET;
    struct sockaddr_in serverAddr;
    HANDLE hThreads[MAX_WORKER_THREADS];

    params.callback = callback;
    
    // 初始化Winsock
    if (WSAStartup(MAKEWORD(2, 2), &wsaData) != 0) {
        fprintf(stderr, "WSAStartup failed\n");
        return;
    }

    // 创建监听套接字
    listenSocket = socket(AF_INET, SOCK_STREAM, IPPROTO_TCP);
    if (listenSocket == INVALID_SOCKET) {
        fprintf(stderr, "Failed to create listen socket, error: %d\n", WSAGetLastError());
        WSACleanup();
        return;
    }

    // 设置套接字为非阻塞模式
    u_long iMode = 1;
    if (ioctlsocket(listenSocket, FIONBIO, &iMode) == SOCKET_ERROR) {
        fprintf(stderr, "ioctlsocket failed with error: %d\n", WSAGetLastError());
        closesocket(listenSocket);
        WSACleanup();
        return;
    }

    serverAddr.sin_family = AF_INET;
    serverAddr.sin_addr.s_addr = INADDR_ANY;
    serverAddr.sin_port = htons(PORT);

    if (bind(listenSocket, (struct sockaddr*)&serverAddr, sizeof(serverAddr)) == SOCKET_ERROR) {
        fprintf(stderr, "Bind failed with error: %d\n", WSAGetLastError());
        closesocket(listenSocket);
        WSACleanup();
        return;
    }

    if (listen(listenSocket, SOMAXCONN) == SOCKET_ERROR) {
        fprintf(stderr, "Listen failed with error: %d\n", WSAGetLastError());
        closesocket(listenSocket);
        WSACleanup();
        return;
    }

    // 创建IO完成端口
    g_hIOCP = CreateIoCompletionPort(INVALID_HANDLE_VALUE, NULL, 0, 0);
    if (g_hIOCP == NULL) {
        fprintf(stderr, "CreateIoCompletionPort failed with error: %d\n", GetLastError());
        closesocket(listenSocket);
        WSACleanup();
        return;
    }
    WorkerThreadParam param;
    // 创建工作者线程
    for (int i = 0; i < MAX_WORKER_THREADS; i++) {
        hThreads[i] = CreateThread(NULL, 0, WorkerThread, NULL, 0, NULL);
        if (hThreads[i] == NULL) {
            fprintf(stderr, "CreateThread failed with error: %d\n", GetLastError());
            g_bShutdown = TRUE;
            break;
        }
    }

    printf("Server started, listening on port %d\n", PORT);

    while (!g_bShutdown) {
        SOCKET clientSocket = accept(listenSocket, NULL, NULL);
        if (clientSocket == INVALID_SOCKET) {
            int error = WSAGetLastError();
            if (error != WSAEWOULDBLOCK) {
                fprintf(stderr, "Accept failed with error: %d\n", error);
                break;
            }
            Sleep(1);  // 避免CPU过度使用
            continue;
        }

        // 设置客户端套接字为非阻塞模式
        u_long iMode = 1;
        if (ioctlsocket(clientSocket, FIONBIO, &iMode) == SOCKET_ERROR) {
            fprintf(stderr, "ioctlsocket failed for client socket with error: %d\n", WSAGetLastError());
            closesocket(clientSocket);
            continue;
        }

        // 分配和初始化 PER_IO_DATA 结构
        LPPER_IO_DATA pIoData = (LPPER_IO_DATA)malloc(sizeof(PER_IO_DATA));
        if (pIoData == NULL) {
            fprintf(stderr, "Failed to allocate memory for PER_IO_DATA\n");
            closesocket(clientSocket);
            continue;
        }
        ZeroMemory(pIoData, sizeof(PER_IO_DATA));
        pIoData->socket = clientSocket;
        pIoData->wsabuf.buf = pIoData->buffer;
        pIoData->wsabuf.len = BUFFER_SIZE;
        pIoData->operation = 0;  // 0 for receive, 1 for send

        // 将客户端套接字与 IO 完成端口关联
        if (CreateIoCompletionPort((HANDLE)clientSocket, g_hIOCP, (ULONG_PTR)pIoData, 0) == NULL) {
            fprintf(stderr, "Failed to associate client socket with IOCP, error: %d\n", GetLastError());
            free(pIoData);
            closesocket(clientSocket);
            continue;
        }

        // 启动异步接收操作
        DWORD flags = 0;
        DWORD bytesReceived = 0;
        if (WSARecv(clientSocket, &pIoData->wsabuf, 1, &bytesReceived, &flags, &pIoData->overlapped, NULL) == SOCKET_ERROR) {
            if (WSAGetLastError() != WSA_IO_PENDING) {
                fprintf(stderr, "WSARecv failed with error: %d\n", WSAGetLastError());
                CleanupSocket(pIoData);
                continue;
            }
        }
    }

    // 清理和关闭
    closesocket(listenSocket);
    CloseHandle(g_hIOCP);
    for (int i = 0; i < MAX_WORKER_THREADS; i++) {
        if (hThreads[i] != NULL) {
            WaitForSingleObject(hThreads[i], INFINITE);
            CloseHandle(hThreads[i]);
        }
    }
    WSACleanup();
}
 
DWORD WINAPI WorkerThread(LPVOID lpParameter) {
    DWORD bytesTransferred;
    LPPER_IO_DATA pIoData;
    LPOVERLAPPED pOverlapped;
    DWORD flags = 0;

    while (!g_bShutdown) {
        BOOL result = GetQueuedCompletionStatus(g_hIOCP, &bytesTransferred, (PULONG_PTR)&pIoData, &pOverlapped, INFINITE);

        if (!result) {
            DWORD error = GetLastError();
            if (error == ERROR_ABANDONED_WAIT_0) {
                break;  // IOCP 句柄被关闭
            }
            fprintf(stderr, "GetQueuedCompletionStatus failed with error: %d\n", error);
            if (pIoData) CleanupSocket(pIoData);
            continue;
        }

        if (bytesTransferred == 0 && pIoData->operation == 0) {
            // 客户端断开连接
            CleanupSocket(pIoData);
            continue;
        }

        if (pIoData->operation == 0) {  // 接收操作
            // 处理接收到的数据
            pIoData->buffer[bytesTransferred] = '\0';
            //printf("Received data: %s\n", pIoData->buffer);

            // 准备发送回显数据
            pIoData->operation = 1;  // 切换到发送操作
            pIoData->wsabuf.len = bytesTransferred;

            ZeroMemory(&pIoData->overlapped, sizeof(OVERLAPPED));
            if (WSASend(pIoData->socket, &pIoData->wsabuf, 1, NULL, 0, &pIoData->overlapped, NULL) == SOCKET_ERROR) {
                if (WSAGetLastError() != WSA_IO_PENDING) {
                    fprintf(stderr, "WSASend failed with error: %d\n", WSAGetLastError());
                    CleanupSocket(pIoData);
                    continue;
                }
            }
        } else {  // 发送操作
            // 发送完成，准备下一次接收
            pIoData->operation = 0;
            pIoData->wsabuf.len = BUFFER_SIZE;
            ZeroMemory(pIoData->buffer, BUFFER_SIZE);
            ZeroMemory(&pIoData->overlapped, sizeof(OVERLAPPED));

            DWORD flags = 0;
            if (WSARecv(pIoData->socket, &pIoData->wsabuf, 1, NULL, &flags, &pIoData->overlapped, NULL) == SOCKET_ERROR) {
                if (WSAGetLastError() != WSA_IO_PENDING) {
                    fprintf(stderr, "WSARecv failed with error: %d\n", WSAGetLastError());
                    CleanupSocket(pIoData);
                    continue;
                }
            }
        }
    }

    return 0;
}

void CleanupSocket(LPPER_IO_DATA pIoData) {
    closesocket(pIoData->socket);
    free(pIoData);
}